# utils/indicators.py
from datetime import timezone
from typing import Optional, List, Dict, Tuple, Any, Callable
import pandas as pd
import pandas_ta as ta  # 确保导入
from entity.k_lines import KLine



class IndicatorCalculator:
    def __init__(self, klines_list: List[KLine]):
        self.df: pd.DataFrame = self._klines_to_dataframe(klines_list)
        self.calculated_indicators: pd.DataFrame = pd.DataFrame(index=self.df.index if not self.df.empty else None)

    def _klines_to_dataframe(self, klines_list: List[KLine]) -> pd.DataFrame:
        # (此方法不变)
        if not klines_list: return pd.DataFrame(columns=['open_time', 'open', 'high', 'low', 'close', 'volume'])
        data = [{'open_time': k.open_time, 'open': k.open, 'high': k.high, 'low': k.low, 'close': k.close,
                 'volume': k.volume,
                 'datetime_pd_utc': pd.to_datetime(k.open_time, unit='ms', utc=True)} for k in klines_list]
        df = pd.DataFrame(data)
        if not df.empty: df.set_index('datetime_pd_utc', inplace=True); df.sort_index(inplace=True)
        return df

    def add_indicator(self,
                      ta_indicator_function: Callable,  # 接收 pandas_ta 的函数本身
                      indicator_name_prefix: str,  # 用于生成列名的前缀
                      source_column: str = 'close',  # 指标计算基于的列，默认为'close'
                      **kwargs) -> None:  # 其他传递给指标函数的参数
        """
        通用的添加指标方法，使用直接函数调用。
        :param ta_indicator_function: pandas_ta库中的指标函数 (例如 ta.ema, ta.rsi)。
        :param indicator_name_prefix: 生成列名的前缀 (例如 "EMA", "RSI")。
        :param source_column: DataFrame中用于计算指标的源数据列名。
        :param kwargs: 传递给pandas_ta指标函数的其他参数 (例如 length, std)。
        """
        if self.df.empty or source_column not in self.df.columns:
            print(f"警告: DataFrame为空或缺少源数据列 '{source_column}'，无法计算指标 {indicator_name_prefix}。")
            return

        try:
            # 直接调用 pandas-ta 函数
            # 例如: indicator_series = ta.ema(close=self.df['close'], length=10)
            indicator_series_or_df = ta_indicator_function(self.df[source_column], **kwargs)

            if indicator_series_or_df is None:
                print(f"警告: 指标 {indicator_name_prefix} 计算返回 None。")
                return

            if isinstance(indicator_series_or_df, pd.Series):
                # 构建列名，例如 EMA_10, RSI_14
                col_name = f"{indicator_name_prefix}"
                if "length" in kwargs: col_name += f"_{kwargs['length']}"
                # 可以根据需要从kwargs中添加更多参数到列名
                if "std" in kwargs and isinstance(kwargs['std'], (int, float)): col_name += f"_{kwargs['std']}"
                if "atr_length" in kwargs: col_name += f"_{kwargs['atr_length']}"  # 例如 Keltner Channels

                self.calculated_indicators[col_name] = indicator_series_or_df
            elif isinstance(indicator_series_or_df, pd.DataFrame):  # 例如 bbands, macd 返回多列
                for col in indicator_series_or_df.columns:
                    # pandas-ta返回的bbands列名可能已经是 BBL_length_stddev 格式
                    # 如果需要，可以在这里调整列名，或者直接使用它返回的
                    self.calculated_indicators[col] = indicator_series_or_df[col]
            else:
                print(f"警告: 指标 {indicator_name_prefix} 计算返回了未知类型: {type(indicator_series_or_df)}")

        except Exception as e:
            print(f"计算指标 {indicator_name_prefix} 时出错: {e}")
            import traceback;
            traceback.print_exc()

    # 便捷方法现在传递函数对象
    def add_ema(self, length: int, column: str = 'close') -> None:
        if column in self.df.columns:
            self.add_indicator(ta.ema, "EMA", source_column=column, length=length)
        else:
            print(f"警告: EMA 计算失败，DataFrame缺少列 '{column}'")

    def add_rsi(self, length: int, column: str = 'close') -> None:
        if column in self.df.columns:
            self.add_indicator(ta.rsi, "RSI", source_column=column, length=length)
        else:
            print(f"警告: RSI 计算失败，DataFrame缺少列 '{column}'")

    def add_bollinger_bands(self, length: int, std: float, column: str = 'close') -> None:
        if column in self.df.columns:
            # bbands 函数会返回DataFrame，add_indicator 会处理多列
            self.add_indicator(ta.bbands, "BBANDS", source_column=column, length=length, std=std)
        else:
            print(f"警告: Bollinger Bands 计算失败，DataFrame缺少列 '{column}'")

    def add_atr(self, length: int) -> None:
        # ATR 需要 high, low, close 列。pandas-ta的 ta.atr 函数签名是:
        # atr(high, low, close, length=None, mamode=None, talib=None, offset=None, **kwargs)
        if not all(col in self.df.columns for col in ['high', 'low', 'close']):
            print("警告: ATR 计算失败，DataFrame缺少 'high', 'low', 或 'close' 列。")
            return
        try:
            atr_series = ta.atr(high=self.df['high'], low=self.df['low'], close=self.df['close'], length=length)
            if atr_series is not None:
                # pandas-ta 返回的ATR列名通常是 "ATRr_length" (例如 ATRr_14)
                # 为了与我们之前的命名一致 (ATR_14)，我们自己指定列名
                self.calculated_indicators[f'ATR_{length}'] = atr_series
        except Exception as e:
            print(f"计算ATR_{length}时出错: {e}")

    # get_kline_and_indicators_at_timestamp 方法保持不变
    def get_kline_and_indicators_at_timestamp(self, timestamp_pd_utc: pd.Timestamp) -> Optional[Dict[str, Any]]:
        # (与上个版本完全相同)
        if self.df.empty: return None
        if timestamp_pd_utc.tzinfo is None:
            timestamp_pd_utc = timestamp_pd_utc.tz_localize('UTC')
        elif timestamp_pd_utc.tzinfo != timezone.utc:
            timestamp_pd_utc = timestamp_pd_utc.astimezone(timezone.utc)
        result = {}
        kline_data = None
        if timestamp_pd_utc in self.df.index:
            kline_data = self.df.loc[timestamp_pd_utc]
        else:
            try:
                idx = self.df.index.get_indexer([timestamp_pd_utc], method='ffill')[0]
                if idx != -1 and self.df.index[idx] <= timestamp_pd_utc: kline_data = self.df.iloc[idx]
            except Exception:
                pass
        if kline_data is not None:
            result.update(kline_data.to_dict())
        else:
            return None
        if not self.calculated_indicators.empty:
            indicators_data = None
            if timestamp_pd_utc in self.calculated_indicators.index:
                indicators_data = self.calculated_indicators.loc[timestamp_pd_utc]
            else:
                try:
                    idx_ind = self.calculated_indicators.index.get_indexer([timestamp_pd_utc], method='ffill')[0]
                    if idx_ind != -1 and self.calculated_indicators.index[idx_ind] <= timestamp_pd_utc:
                        indicators_data = self.calculated_indicators.iloc[idx_ind]
                except Exception:
                    pass
            if indicators_data is not None: result.update(indicators_data.to_dict())
        return result if result else None